maybe we shall also pass the session object here?

Do you mean move the activeSessions.remove(session) call into this function like so?

private def handleSessionTermination(
     session: WorkerSession,
     workerSpec: UDFWorkerSpecification
  ): Unit = {
    activeSessions.remove(session)
    
    val entry = dispatchers.get(workerSpec)
    // Note: entry == null is unexpected and should
    // throw here.
    entry.activeSessionCount -= 1
    if (entry.activeSessionCount == 0) {
      logger.info("All sessions closed for dispatcher " +
        s"${entry.dispatcher.dispatcherId}, removing from cache")
      dispatchers.remove(workerSpec)
      onAllDispatcherSessionsClosed(entry.dispatcher)
    }
  }

it may be better to include this per dispatcher, rather in this factory.

In my view the point of this function is to decide how long a dispatcher is being kept around. E.g. in a scenario where clients provide different worker specs, we cannot keep every dispatcher for the whole Spark runtime. There needs to be some sort of timeout, especially once we implement warm-pooling.

Should this dispatcher timeout be controlled by each dispatcher instead of a central place? The re-use logic may also depend on how many dispatchers there currently are. E.g. if we only need one dispatcher in the whole Spark lifetime there is no need for eager cleanup. Therefore, I would argue its more suitable to decide this in the factory, which has knowledge of the number and kind of dispatchers that are currently in existence

same as this - factory pattern means it only defines the object creation logic, onStop should be added to dispatcher instance.

These logic looks like more dispatcher impl. details - e.g. a proper implemented dispatcher should manage active sessions properly.

A dispatcher spawns workers and creates sessions, then it's more natural a dispatcher manages the sessions it starts.

For example, the dispatcher manager here creates dispatcher, then it takes care of disposal of dispatchers.

A dispatcher spawns workers and sessions, then the disposal of session and workers should live in a dispatcher.

Because creation and disposal logic can be related, it's better to colocate them together.